# This Python file uses the following encoding: utf-8
import sys
import socket
import random
from PySide6.QtWidgets import (
    QApplication, QMainWindow, QWidget, QLabel, QLineEdit, QTextEdit, QPushButton, QVBoxLayout, QHBoxLayout, QTabWidget, QCheckBox, QSplitter, QFileDialog, QSizePolicy
)
from PySide6.QtCore import QTimer, Signal, Qt
# from PySide6.QtGui import QTextCursor
import threading
import datetime
import re

class TcpWidget(QWidget):
    startExposureTimer = Signal(int)
    log_signal = Signal(str)
    first_recv_signal = Signal(str)
    append_recv_signal = Signal(str)

    def __init__(self, udp_widget=None):
        super().__init__()
        self.udp_widget = udp_widget
        ip_label = QLabel("目标IP地址:")
        self.ip_input = QLineEdit()
        self.ip_input.setFixedWidth(140)
        port_label = QLabel("目标端口号:")
        self.port_input = QLineEdit()
        self.port_input.setFixedWidth(90)
        local_port_label = QLabel("本地端口:")
        local_port_label.setSizePolicy(QSizePolicy.Minimum, QSizePolicy.Fixed)
        self.local_port_input = QLineEdit()
        self.local_port_input.setFixedWidth(90)
        self.local_port_input.setText("8080")

        self.connect_button = QPushButton("连接")
        self.connect_button.setFixedWidth(60)
        self.connect_button.setFixedHeight(28)
        self.connect_button.clicked.connect(self.connect_to_server)
        self.disconnect_button = QPushButton("断开")
        self.disconnect_button.setFixedWidth(60)
        self.disconnect_button.setFixedHeight(28)
        self.disconnect_button.clicked.connect(self.disconnect_from_server)
        self.disconnect_button.setEnabled(False)

        # 顶部设置区两行
        top_row1 = QHBoxLayout()
        top_row1.addWidget(ip_label)
        top_row1.addWidget(self.ip_input)
        top_row1.addWidget(port_label)
        top_row1.addWidget(self.port_input)
        top_row1.addStretch(1)

        top_row2 = QHBoxLayout()
        top_row2.addWidget(local_port_label)
        top_row2.addWidget(self.local_port_input)
        top_row2.addWidget(self.connect_button)
        top_row2.addWidget(self.disconnect_button)
        top_row2.addStretch(1)

        # 曝光指令模块
        exposure_group = QWidget()
        exposure_layout = QVBoxLayout()
        exposure_label = QLabel("曝光指令:")
        self.exposure_send_input = QTextEdit()
        # self.exposure_send_input.setMinimumWidth(350)  # 删除或注释
        self.exposure_send_input.setLineWrapMode(QTextEdit.WidgetWidth)
        self.exposure_send_input.setSizePolicy(QSizePolicy.Expanding, QSizePolicy.Fixed)
        # 读出时间输入框
        readout_time_layout = QHBoxLayout()
        self.readout_time_label = QLabel("读出时间(s):")
        self.readout_time_input = QLineEdit()
        self.readout_time_input.setText("10")
        self.readout_time_input.setFixedWidth(60)
        readout_time_layout.addWidget(self.readout_time_label)
        readout_time_layout.addWidget(self.readout_time_input)
        # 发送曝光指令按钮
        self.exposure_send_button = QPushButton("发送曝光指令")
        self.exposure_send_button.setMinimumHeight(32)
        self.exposure_send_button.clicked.connect(self.send_exposure_message)
        # 批量曝光相关控件
        self.batch_file_button = QPushButton("选择批量文件")
        self.batch_file_button.setMinimumHeight(32)
        self.batch_file_button.clicked.connect(self.select_batch_file)
        self.batch_file_path = QLineEdit()
        self.batch_file_path.setMinimumHeight(32)
        self.batch_file_path.setReadOnly(True)
        self.batch_start_button = QPushButton("批量曝光")
        self.batch_start_button.setMinimumHeight(32)
        self.batch_start_button.clicked.connect(self.start_batch_exposure)
        self.batch_start_button.setEnabled(False)
        batch_btn_layout = QHBoxLayout()
        batch_btn_layout.addWidget(self.batch_file_button)
        batch_btn_layout.addWidget(self.batch_start_button)
        batch_path_layout = QHBoxLayout()
        batch_path_layout.addWidget(self.batch_file_path)
        exposure_layout.addWidget(exposure_label)
        exposure_layout.addWidget(self.exposure_send_input)
        exposure_layout.addLayout(readout_time_layout)
        exposure_layout.addWidget(self.exposure_send_button)
        exposure_layout.addLayout(batch_btn_layout)
        exposure_layout.addLayout(batch_path_layout)
        exposure_group.setLayout(exposure_layout)

        # 其他指令模块
        normal_group = QWidget()
        normal_layout = QVBoxLayout()
        normal_label = QLabel("其他指令:")
        self.normal_send_input = QTextEdit()
        # self.normal_send_input.setMinimumWidth(350)  # 删除或注释
        self.normal_send_input.setLineWrapMode(QTextEdit.WidgetWidth)
        self.normal_send_input.setSizePolicy(QSizePolicy.Expanding, QSizePolicy.Fixed)
        self.normal_send_button = QPushButton("发送指令")
        self.normal_send_button.setMinimumHeight(32)
        self.normal_send_button.clicked.connect(self.send_normal_message)
        normal_layout.addWidget(normal_label)
        normal_layout.addWidget(self.normal_send_input)
        normal_layout.addWidget(self.normal_send_button)
        normal_layout.addSpacing(self.exposure_send_button.sizeHint().height() + batch_btn_layout.sizeHint().height() + batch_path_layout.sizeHint().height())
        normal_group.setLayout(normal_layout)

        # 指令发送区域左右布局
        send_split_layout = QHBoxLayout()
        send_split_layout.addWidget(exposure_group, 1)
        send_split_layout.addWidget(normal_group, 1)

        # 返回信息框
        recv_label = QLabel("返回信息:")
        self.recv_output = QTextEdit()
        self.recv_output.setReadOnly(True)
        # 全局清除返回信息按钮
        self.clear_button = QPushButton("清除返回信息")
        self.clear_button.setMinimumHeight(32)
        self.clear_button.clicked.connect(self.clear_recv_output)

        # 持续连接相关
        self.tcp_socket = None
        self.listen_thread = None
        self.listening = False
        self.exposure_state = "idle"
        # 移除定时器相关代码，不再需要等待时间
        self.last_exposure_command = None
        # 批量状态
        self.batch_commands = []
        self.batch_index = 0
        self.batch_running = False
        self.waiting_first_response = False
        self.first_recv_signal.connect(self.recv_output.setPlainText)
        self.append_recv_signal.connect(self.recv_output.append)

        # layout
        main_layout = QVBoxLayout()
        main_layout.addLayout(top_row1)
        main_layout.addLayout(top_row2)
        main_layout.addLayout(send_split_layout)
        main_layout.addWidget(recv_label)
        main_layout.addWidget(self.recv_output)
        main_layout.addWidget(self.clear_button)
        self.setLayout(main_layout)
        # 设置窗口初始更宽
        self.setMinimumWidth(800)
        self.setMinimumHeight(600)

    def connect_to_server(self):
        ip = self.ip_input.text().strip()
        port_text = self.port_input.text().strip()
        local_port_text = self.local_port_input.text().strip()
        if not ip or not port_text or not local_port_text:
            self.append_recv_signal.emit("请填写完整的目标IP、端口和本地端口！")
            return
        try:
            port = int(port_text)
            local_port = int(local_port_text)
        except ValueError:
            self.append_recv_signal.emit("端口号必须为整数！")
            return
        try:
            self.tcp_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            self.tcp_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
            self.tcp_socket.bind(("0.0.0.0", local_port))
            self.tcp_socket.connect((ip, port))
            self.listening = True
            self.listen_thread = threading.Thread(target=self.listen_loop, daemon=True)
            self.listen_thread.start()
            self.connect_button.setEnabled(False)
            self.disconnect_button.setEnabled(True)
            self.append_recv_signal.emit(f"已连接到 {ip}:{port}，本地端口 {local_port}")
        except Exception as e:
            self.append_recv_signal.emit(f"连接失败: {e}")
            self.tcp_socket = None

    def disconnect_from_server(self):
        self.listening = False
        if self.tcp_socket:
            try:
                self.tcp_socket.shutdown(socket.SHUT_RDWR)
            except:
                pass
            try:
                self.tcp_socket.close()
            except:
                pass
            self.tcp_socket = None
        self.connect_button.setEnabled(True)
        self.disconnect_button.setEnabled(False)
        self.append_recv_signal.emit("已断开连接")

    def listen_loop(self):
        while self.listening and self.tcp_socket:
            try:
                data = self.tcp_socket.recv(4096)
                if not data:
                    self.first_recv_signal.emit("连接已关闭")
                    self.log("连接已关闭")
                    self.disconnect_from_server()
                    break
                msg = data.decode(errors='ignore')
                if self.waiting_first_response:
                    self.first_recv_signal.emit(f"收到: {msg}")
                    self.waiting_first_response = False
                self.log(f"收到: {msg}")
                # 曝光流程自动处理：收到OK响应后直接等待完成信号
                if self.exposure_state == "waiting_first_response" and "<INFO=OK>" in msg and self.last_exposure_command:
                    self.exposure_state = "waiting_second_response"
                    self.last_exposure_command = None
                # 监听到包含<ID=302>和<ACTION=DONE>的响应后，批量模式自动进入下一条
                if self.batch_running and "<ID=302>" in msg and "<ACTION=DONE>" in msg:
                    self.batch_index += 1
                    self.run_next_batch_command()
            except Exception as e:
                if self.listening:
                    self.first_recv_signal.emit(f"监听异常: {e}")
                    self.log(f"监听异常: {e}")
                # 批量模式下遇到异常也进入下一条
                if self.batch_running:
                    self.batch_index += 1
                    QTimer.singleShot(500, self.run_next_batch_command)
                break

    def send_exposure_message(self):
        if not self.tcp_socket:
            self.append_recv_signal.emit("未连接，无法发送！")
            return
        message = self.exposure_send_input.toPlainText()
        if not message:
            self.append_recv_signal.emit("请输入曝光指令内容！")
            return
        self.recv_output.clear()
        self.waiting_first_response = True
        try:
            self.tcp_socket.sendall(message.encode())
            self.append_recv_signal.emit(f"发送曝光指令: {message}")
            self.log(f"发送曝光指令: {message}")
            self.exposure_state = "waiting_first_response"
            self.last_exposure_command = message
        except Exception as e:
            self.append_recv_signal.emit(f"发送失败: {e}")
            self.log(f"发送失败: {e}")

    def send_normal_message(self):
        if not self.tcp_socket:
            self.append_recv_signal.emit("未连接，无法发送！")
            return
        message = self.normal_send_input.toPlainText()
        if not message:
            self.append_recv_signal.emit("请输入普通指令内容！")
            return
        self.recv_output.clear()
        self.waiting_first_response = True
        try:
            self.tcp_socket.sendall(message.encode())
            self.append_recv_signal.emit(f"发送普通指令: {message}")
            self.log(f"发送普通指令: {message}")
        except Exception as e:
            self.append_recv_signal.emit(f"发送失败: {e}")
            self.log(f"发送失败: {e}")

    def send_read_image_command(self):
        # 只负责自动发送UDP包，不做UI端口校验
        if not self.udp_widget:
            self.append_recv_signal.emit("未配置UDP通信端口，无法发送读出指令！")
            self.log("未配置UDP通信端口，无法发送读出指令！")
            return
        ip = self.udp_widget.ip_input.text().strip()
        port_text = self.udp_widget.port_input.text().strip()
        try:
            port = int(port_text)
        except Exception:
            self.append_recv_signal.emit("UDP端口号无效！")
            self.log("UDP端口号无效！")
            return
        message_to_send = "get_frame_290:done=01"
        try:
            with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as s:
                s.sendto(message_to_send.encode(), (ip, port))
            self.append_recv_signal.emit(f"已通过UDP发送读出图像指令: {message_to_send} 到 {ip}:{port}")
            self.log(f"已通过UDP发送读出图像指令: {message_to_send} 到 {ip}:{port}")
            # 移除这里的自动推进，让listen_loop中的检测逻辑来处理
        except Exception as e:
            self.append_recv_signal.emit(f"UDP发送失败: {e}")
            self.log(f"UDP发送失败: {e}")

    def clear_recv_output(self):
        self.recv_output.clear()
        self.exposure_state = "idle"

    def closeEvent(self, event):
        self.listening = False
        if self.tcp_socket:
            try:
                self.tcp_socket.shutdown(socket.SHUT_RDWR)
            except:
                pass
            try:
                self.tcp_socket.close()
            except:
                pass
            self.tcp_socket = None
        event.accept()

    def log(self, msg):
        timestamp = datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')
        self.log_signal.emit(f"[{timestamp}] {msg}")

    def select_batch_file(self):
        file_path, _ = QFileDialog.getOpenFileName(self, "选择批量曝光指令文件", "", "Text Files (*.txt)")
        if file_path:
            self.batch_file_path.setText(file_path)
            self.batch_start_button.setEnabled(True)
        else:
            self.batch_file_path.clear()
            self.batch_start_button.setEnabled(False)

    def start_batch_exposure(self):
        file_path = self.batch_file_path.text().strip()
        if not file_path:
            self.append_recv_signal.emit("请先选择批量曝光指令文件！")
            return
        try:
            with open(file_path, 'r', encoding='utf-8') as f:
                lines = [line.strip() for line in f if line.strip()]
            if not lines:
                self.append_recv_signal.emit("文件内容为空！")
                return
            self.batch_commands = lines
            self.batch_index = 0
            self.batch_running = True
            self.append_recv_signal.emit(f"开始批量曝光，共{len(lines)}条指令...")
            self.log(f"开始批量曝光，共{len(lines)}条指令...")
            self.run_next_batch_command()
        except Exception as e:
            self.append_recv_signal.emit(f"读取文件失败: {e}")
            self.batch_running = False

    def run_next_batch_command(self):
        if not self.batch_running or self.batch_index >= len(self.batch_commands):
            self.append_recv_signal.emit("批量曝光完成！")
            self.log("批量曝光完成！")
            self.batch_running = False
            return
        # 设置曝光指令并发送
        cmd = self.batch_commands[self.batch_index]
        # 不再在曝光指令框显示内容，直接发送
        if self.tcp_socket:
            self.waiting_first_response = True
            try:
                self.tcp_socket.sendall(cmd.encode())
                self.append_recv_signal.emit(f"发送曝光指令: {cmd}")
                self.log(f"发送曝光指令: {cmd}")
                self.exposure_state = "waiting_first_response"
                self.last_exposure_command = cmd
            except Exception as e:
                self.append_recv_signal.emit(f"发送失败: {e}")
                self.log(f"发送失败: {e}")
        else:
            self.append_recv_signal.emit("未连接，无法发送！")
            self.log("未连接，无法发送！")
        # send_message会自动进入曝光流程，流程结束后需自动调用本函数

class UdpWidget(QWidget):
    log_signal = Signal(str)
    first_recv_signal = Signal(str)
    append_recv_signal = Signal(str)
    setplain_signal = Signal(str)

    def __init__(self):
        super().__init__()
        
        ip_label = QLabel("目标IP地址:")
        self.ip_input = QLineEdit()
        
        port_label = QLabel("目标端口号:")
        self.port_input = QLineEdit()
        
        local_port_label = QLabel("本地监听端口:")
        self.local_port_input = QLineEdit()
        self.local_port_input.setText("8001")  # 默认本地端口8001
        self.local_port_status = QLabel("当前监听端口: 8001")
        self.local_port_input.textChanged.connect(self.update_local_port_status)
        
        send_label = QLabel("发送内容:")
        self.send_input = QTextEdit()
        
        self.send_button = QPushButton("发送")
        self.send_button.clicked.connect(self.send_message)
        
        self.clear_button = QPushButton("清除返回信息")
        self.clear_button.clicked.connect(self.clear_recv_output)
        
        recv_label = QLabel("返回信息:")
        self.recv_output = QTextEdit()
        self.recv_output.setReadOnly(True)
        # layout
        ip_port_layout = QHBoxLayout()
        ip_port_layout.addWidget(ip_label)
        ip_port_layout.addWidget(self.ip_input)
        ip_port_layout.addWidget(port_label)
        ip_port_layout.addWidget(self.port_input)
        local_port_layout = QHBoxLayout()
        local_port_layout.addWidget(local_port_label)
        local_port_layout.addWidget(self.local_port_input)
        local_port_layout.addWidget(self.local_port_status)
        main_layout = QVBoxLayout()
        main_layout.addLayout(ip_port_layout)
        main_layout.addLayout(local_port_layout)
        main_layout.addWidget(send_label)
        main_layout.addWidget(self.send_input)
        main_layout.addWidget(self.send_button)
        main_layout.addWidget(self.clear_button)
        main_layout.addWidget(recv_label)
        main_layout.addWidget(self.recv_output)
        self.setLayout(main_layout)
        # 监听线程
        self.listen_thread = None
        self.listening_port = None
        self.update_local_port_status()  # 启动监听
        self.waiting_first_response = False
        self.first_recv_signal.connect(self.recv_output.setPlainText)
        self.append_recv_signal.connect(self.recv_output.append)
        self.setplain_signal.connect(self.recv_output.setPlainText)

    def update_local_port_status(self):
        port = self.local_port_input.text().strip()
        if port:
            self.local_port_status.setText(f"当前监听端口: {port}")
        else:
            self.local_port_status.setText("当前监听端口: 未设置")
        self.restart_listen_thread()

    def restart_listen_thread(self):
        port_text = self.local_port_input.text().strip()
        try:
            port = int(port_text)
            if port <= 0 or port > 65535:
                raise ValueError
        except Exception:
            port = None
        
        if port == self.listening_port and self.listen_thread and self.listen_thread.is_alive():
            return
        if self.listen_thread:
            self.stop_listen_thread()
        if port:
            import threading
            self.listening_port = port
            self.listen_thread_running = True
            self.listen_thread = threading.Thread(target=self.listen_loop, args=(port,), daemon=True)
            self.listen_thread.start()
        else:
            self.listening_port = None

    def stop_listen_thread(self):
        self.listen_thread_running = False
        import time
        time.sleep(0.1)
        self.listen_thread = None

    def listen_loop(self, port):
        with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as s:
            s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
            try:
                s.bind(("0.0.0.0", port))
            except Exception as e:
                self.append_recv_signal.emit(f"监听端口失败: {e}")
                return
            self.append_recv_signal.emit(f"已开始监听端口: {port}")
            while self.listen_thread_running:
                try:
                    s.settimeout(1)
                    data, addr = s.recvfrom(4096)
                    msg = data.decode(errors='ignore')
                    if self.waiting_first_response:
                        self.first_recv_signal.emit(f"{msg}")
                        self.waiting_first_response = False
                    self.log(f"收到: {msg}")
                    # 收到指定信息后的自动回复逻辑，包括temp_read、vaccu_read
                    if msg.strip() == 'temp_read':
                        arr = [random.uniform(-110, -80) for _ in range(12)]
                        arr += [random.uniform(60, 80) for _ in range(2)]
                        arr.append(random.uniform(-110, -80))
                        self.log(f"temp: {arr}")
                        fields = [
                            'CCD_3', 'CCD_2', 'Base_1', 'Absorbent_cold', 'Drive_circuit', 'Backup1',
                            'CCD_4', 'CCD_1', 'Base_2', 'Absorbent_hot', 'Dewar', 'Backup2',
                            'output', 'output_ab', 'temp_set'
                        ]
                        reply = '&&'.join([f"{name}:{value:.2f}" for name, value in zip(fields, arr)])
                        s.sendto(reply.encode(), addr)
                    elif msg.strip() == 'vaccu_read':
                        prefix = '0011074006'
                        # 前4位有效数字
                        middle_first = f"{random.randint(50, 2100):04d}"
                        # 后2位指数幂
                        middle_second = f"{random.randint(5, 18):02d}"
                        middle = middle_first + middle_second
                        suffix = '038'
                        vaccu_reply ='vaccu_read:' + prefix + middle + suffix
                        self.log(f"vaccu: {prefix + middle + suffix}")
                        s.sendto(vaccu_reply.encode(), addr)
                    else:
                        reply = '已收到信息'
                        # self.log(f"收到信息: {msg}")
                        s.sendto(reply.encode(), addr)
                except socket.timeout:
                    continue
                except Exception as e:
                    self.append_recv_signal.emit(f"监听异常: {e}")
                    self.log(f"监听异常: {e}")
                    break
            self.append_recv_signal.emit("监听已停止")

    def send_message(self):
        ip = self.ip_input.text().strip()
        port_text = self.port_input.text().strip()
        message = self.send_input.toPlainText()
        if not ip or not port_text or not message:
            self.setplain_signal.emit("请填写完整的目标IP、端口和发送内容！")
            return
        self.recv_output.clear()
        self.waiting_first_response = True
        try:
            port = int(port_text)
        except ValueError:
            self.setplain_signal.emit("请检查IP和端口号")
            return
        try:
            with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as s:
                s.settimeout(2)  # UDP通信线程持续时间
                s.sendto(message.encode(), (ip, port))
                self.log(f"发送: {message}")
                try:
                    data, _ = s.recvfrom(4096)
                    recv_msg = data.decode(errors='ignore')
                    if self.waiting_first_response:
                        self.setplain_signal.emit(recv_msg)
                        self.waiting_first_response = False
                    # self.log(f"收到: {recv_msg}")
                except socket.timeout:
                    self.setplain_signal.emit("接收超时，无返回数据")
                    self.log("接收超时，无返回数据")
        except Exception as e:
            self.setplain_signal.emit(f"错误: {e}")
            self.log(f"错误: {e}")

    def clear_recv_output(self):
        self.recv_output.clear()

    def closeEvent(self, event):
        self.stop_listen_thread()
        event.accept()

    def log(self, msg):
        timestamp = datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')
        self.log_signal.emit(f"[{timestamp}] {msg}")

class MainWindow(QMainWindow):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("ICS测试工具")
        self.resize(1100, 600)
        self.tabs = QTabWidget()
        self.udp_widget = UdpWidget()
        self.tcp_widget = TcpWidget(self.udp_widget)
        self.tabs.addTab(self.tcp_widget, "OCS模拟(TCP)")
        self.tabs.addTab(self.udp_widget, "温控板模拟(UDP)")
        # 日志窗口
        self.log_output = QTextEdit()
        self.log_output.setReadOnly(True)
        self.log_output.setMinimumWidth(350)
        splitter = QSplitter(Qt.Horizontal)
        splitter.addWidget(self.tabs)
        splitter.addWidget(self.log_output)
        splitter.setSizes([600, 600])  # 让日志区和左侧tab页一样宽
        self.setCentralWidget(splitter)
        # 信号连接
        self.tcp_widget.log_signal.connect(self.append_log)
        self.udp_widget.log_signal.connect(self.append_log)

    def append_log(self, msg):
        self.log_output.append(msg)

if __name__ == "__main__":
    app = QApplication(sys.argv)
    window = MainWindow()
    window.show()
    sys.exit(app.exec())
